/**
 * @file laserOdometry.hpp
 * @author Linfu Wei (ghowoght@qq.com)
 * @brief
 * @version 1.0
 * @date 2022-05-06
 *
 * @copyright Copyright (c) 2022
 *
 */

#include <fstream>

#include <vector>
#include <random>
#include <future>
#include <map>
#include <queue>
#include <mutex>
#include <thread>

#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <nav_msgs/Path.h>
#include <tf/transform_broadcaster.h>

#include <pcl/point_types.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/common/centroid.h>
#include <pcl/common/eigen.h>

#include <Eigen/Core>
#include <Eigen/Eigenvalues>

#include <ceres/ceres.h>

#include "lidarFactor.hpp"
#include "tictoc.hpp"

#include "ikd-Tree/ikd_Tree.h"

using PointType = pcl::PointXYZINormal;
using PointCloud = pcl::PointCloud<PointType>;
using PointVector = KD_TREE<PointType>::PointVector;
using V3D = Eigen::Vector3d;
using M3D = Eigen::Matrix3d;
using QD = Eigen::Quaterniond;
template class KD_TREE<PointType>;

namespace simple_loam
{


    // typedef pcl::PointXYZINormal PointType;
    typedef pcl::PointCloud<PointType> PointCloudXYZI;

    struct DataPair
    {
        sensor_msgs::PointCloud2ConstPtr cloud_edge;
        sensor_msgs::PointCloud2ConstPtr cloud_surf;
        sensor_msgs::PointCloud2ConstPtr cloud_full;
        bool has_edge = false;
        bool has_surf = false;
        bool has_full = false;
        void add_cloud_edge(sensor_msgs::PointCloud2ConstPtr cloud)
        {
            cloud_edge = cloud;
            has_edge = true;
        }
        void add_cloud_surf(sensor_msgs::PointCloud2ConstPtr cloud)
        {
            cloud_surf = cloud;
            has_surf = true;
        }
        void add_cloud_full(sensor_msgs::PointCloud2ConstPtr cloud)
        {
            cloud_full = cloud;
            has_full = true;
        }
        bool is_completed()
        {
            return has_edge && has_surf && has_full;
        }
    };

    class LaserOdometry
    {

    private:
        ros::NodeHandle nh_;
        ros::Subscriber sub_edge_cloud;
        ros::Subscriber sub_surf_cloud;
        ros::Subscriber sub_full_cloud;
        ros::Publisher pub_surf_map;
        ros::Publisher pub_edge_map;
        ros::Publisher pub_traj;

        tf::TransformBroadcaster odom_broadcaster;
        geometry_msgs::TransformStamped odom_trans;
        nav_msgs::Path traj_local_msgs;

        bool is_initial = false;
        double time_edge_cloud;
        double time_surf_cloud;
        std::string frame_id;
        PointCloudXYZI::Ptr cloud_edge_curr; // current edge cloud
        PointCloudXYZI::Ptr cloud_surf_curr; // current surf cloud
        PointCloudXYZI::Ptr map_edge; // edge map
        PointCloudXYZI::Ptr map_surf; // surf map
        std::vector<PointCloudXYZI::Ptr> map_edge_vec; // edge map vector
        std::vector<PointCloudXYZI::Ptr> map_surf_vec; // surf map vector

        pcl::KdTreeFLANN<PointType> kdtree_edge_map; // kdtree for edge map
        pcl::KdTreeFLANN<PointType> kdtree_surf_map; // kdtree for surf map

        KD_TREE<PointType>::Ptr ikd_tree_;
        int KNN_MATCH_NUM = 8; // K近邻匹配点数
        double map_res_ = 0.5;
        PointVector nearest;    // 最近邻搜索结果
        std::vector<PointVector> nearest_queue_;

        std::map<double, DataPair *> map_data_pair;
        std::queue<DataPair *> queue_avail_data;
        std::mutex mutex_queue_avail_data;


        double q_params_increm[4] = {0, 0, 0, 1}; // ceres待优化参数
        double t_params_increm[3] = {0, 0, 0};    // ceres待优化参数

        std::ofstream result_fs;

    private:
        DataPair *get_data_pair(double timestamp)
        {
            std::map<double, DataPair *>::iterator it = map_data_pair.find(timestamp);
            if (it == map_data_pair.end())
            { // 如果没有，则创建一个新的键值对
                DataPair *data_pair = new DataPair();
                map_data_pair.insert(std::make_pair(timestamp, data_pair));
                return data_pair;
            }
            else
                return it->second; // first: 返回key | second: 返回value
        }

    private:
        void edgeCloudHandler(const sensor_msgs::PointCloud2ConstPtr &msg)
        {
            std::unique_lock<std::mutex> lock(mutex_queue_avail_data);
            auto data_pair = get_data_pair(msg->header.stamp.toSec());
            data_pair->add_cloud_edge(msg);
            if (data_pair->is_completed())
            {
                queue_avail_data.push(data_pair);
            }
        }

        void surfCloudHandler(const sensor_msgs::PointCloud2ConstPtr &msg)
        {
            std::unique_lock<std::mutex> lock(mutex_queue_avail_data);
            auto data_pair = get_data_pair(msg->header.stamp.toSec());
            data_pair->add_cloud_surf(msg);
            if (data_pair->is_completed())
            {
                queue_avail_data.push(data_pair);
            }
        }
        void fullCloudHandler(const sensor_msgs::PointCloud2ConstPtr &msg)
        {
            std::unique_lock<std::mutex> lock(mutex_queue_avail_data);
            auto data_pair = get_data_pair(msg->header.stamp.toSec());
            data_pair->add_cloud_full(msg);
            if (data_pair->is_completed())
            {
                queue_avail_data.push(data_pair);
            }
        }

        void map_update(PointCloud::Ptr& features){
            TicToc t_build_map;
            PointCloud::Ptr points_to_add = boost::make_shared<PointCloud>();
            PointCloud::Ptr points_not_to_downsample = boost::make_shared<PointCloud>();
            
            for(int i = 0; i < features->size(); i++){
                auto& nearest = nearest_queue_[i];
                V3D pb((*features)[i].x, (*features)[i].y, (*features)[i].z);
                V3D pw = q_w_curr * pb + t_w_curr; // 转换到世界坐标系下

                if(pb.norm() > 150) continue;

                PointType p;
                p.x = pw.x();
                p.y = pw.y();
                p.z = pw.z();

                PointType mid_point;
                mid_point.x = floor(pw.x() / map_res_) * map_res_ + map_res_ / 2;
                mid_point.y = floor(pw.y() / map_res_) * map_res_ + map_res_ / 2;
                mid_point.z = floor(pw.z() / map_res_) * map_res_ + map_res_ / 2;
                V3D p_mid(mid_point.x, mid_point.y, mid_point.z);
                double dist = (pw - p_mid).norm();
                if(fabs(nearest[0].x - mid_point.x) > map_res_ / 2
                    && fabs(nearest[0].y - mid_point.y) > map_res_ / 2
                    && fabs(nearest[0].z - mid_point.z) > map_res_ / 2){
                    points_not_to_downsample->push_back(p);
                    continue;
                }
                bool need_add = true;
                for(int k = 0; k < KNN_MATCH_NUM; k++){
                    if(nearest.size() < KNN_MATCH_NUM)
                        break;
                    V3D pn(nearest[k].x, nearest[k].y, nearest[k].z);
                    if((pn - p_mid).norm() < dist){
                        need_add = false;
                        break;
                    }
                }
                if(need_add){
                    points_to_add->push_back(p);
                }
            }
            ikd_tree_->Add_Points(points_to_add->points, true);
            ikd_tree_->Add_Points(points_not_to_downsample->points, false);
        }

    public:
        LaserOdometry(ros::NodeHandle &nh) : nh_(nh)
        {

            result_fs.open("/home/ghowoght/workspace/lidar_ws/src/simple_loam/scripts/result.txt");

            map_edge = boost::make_shared<PointCloudXYZI>();
            map_surf = boost::make_shared<PointCloudXYZI>();

            cloud_surf_curr = boost::make_shared<PointCloudXYZI>();
            cloud_edge_curr = boost::make_shared<PointCloudXYZI>();

            sub_edge_cloud = nh_.subscribe<sensor_msgs::PointCloud2>("/laser_edge_cloud", 100,
                                                                     boost::bind(&LaserOdometry::edgeCloudHandler, this, _1));
            sub_surf_cloud = nh_.subscribe<sensor_msgs::PointCloud2>("/laser_surf_cloud", 100,
                                                                     boost::bind(&LaserOdometry::surfCloudHandler, this, _1));
            sub_full_cloud = nh_.subscribe<sensor_msgs::PointCloud2>("/laser_cloud", 100,
                                                                     boost::bind(&LaserOdometry::fullCloudHandler, this, _1));

            pub_edge_map = nh_.advertise<sensor_msgs::PointCloud2>("/map/edge", 10);
            pub_surf_map = nh_.advertise<sensor_msgs::PointCloud2>("/map/surf", 10);
            pub_traj     = nh_.advertise<nav_msgs::Path>("/traj/local", 10);

            ikd_tree_ = std::make_shared<KD_TREE<PointType>>(); // delete_param, balance_param, box_length
            ikd_tree_->set_downsample_param(map_res_);

            map_surf_cloud_ = boost::make_shared<PointCloudXYZI>();

            std::future<void> thd(std::async(std::launch::async, &simple_loam::LaserOdometry::spin, this));
            ros::spin();
        }

        Eigen::Quaterniond q_w_curr; // b->w的旋转四元数
        Eigen::Vector3d t_w_curr; // b->w的平移向量
        PointCloud::Ptr map_surf_cloud_;
        void spin()
        {

            Eigen::Map<Eigen::Quaterniond> q_w_increm(q_params_increm); // b->w的旋转四元数
            Eigen::Map<Eigen::Vector3d> t_w_increm(t_params_increm);    // b->w的平移向量

            // Eigen::Quaterniond q_w_curr(0, 0, 0, 1); // b->w的旋转四元数
            // Eigen::Vector3d t_w_curr(0, 0, 0);       // b->w的平移向量
            q_w_curr = Eigen::Quaterniond(1, 0, 0, 0);
            t_w_curr = Eigen::Vector3d(0, 0, 0);
            Eigen::Quaterniond q_w_last = q_w_curr;
            Eigen::Vector3d t_w_last = t_w_curr;

            while (ros::ok())
            {

                std::this_thread::sleep_for(std::chrono::nanoseconds(1));

                if (!queue_avail_data.empty())
                {
                    // 读取当前帧edge和surf点云数据
                    auto data_pair = queue_avail_data.front();
                    queue_avail_data.pop();
                    time_surf_cloud = data_pair->cloud_surf->header.stamp.toSec();
                    pcl::fromROSMsg(*data_pair->cloud_surf, *cloud_surf_curr);
                    time_edge_cloud = data_pair->cloud_edge->header.stamp.toSec();
                    pcl::fromROSMsg(*data_pair->cloud_edge, *cloud_edge_curr);

                    frame_id = data_pair->cloud_edge->header.frame_id;

                    delete data_pair;

                    TicToc t_total;
                    if (!is_initial)
                    { 
                        // 初始化点云地图                       
                        *map_surf_cloud_ += *cloud_surf_curr;

                        PointCloudXYZI::Ptr surf_filtered(new PointCloudXYZI);
                        pcl::VoxelGrid<PointType> sor;
                        sor.setInputCloud(map_surf_cloud_);
                        sor.setLeafSize(map_res_, map_res_, map_res_);
                        sor.filter(*surf_filtered);
                        map_surf_cloud_ = surf_filtered;

                        if(map_surf_cloud_->size() > 1000){
                            is_initial = true;

                            ikd_tree_->Build(map_surf_cloud_->points);
                        }
                    }
                    else
                    {
                        // 位姿增量清零
                        q_params_increm[0] = q_params_increm[1] = q_params_increm[2] = 0;
                        q_params_increm[3] = 1;
                        t_params_increm[0] = t_params_increm[1] = t_params_increm[2] = 0;

                        // 对当前帧平面点降采样
                        if (1)
                        {
                            pcl::VoxelGrid<PointType> surf_filter;
                            surf_filter.setLeafSize(0.2, 0.2, 0.2);
                            surf_filter.setInputCloud(cloud_surf_curr);
                            PointCloudXYZI::Ptr cloud_surf_temp_filtered(new PointCloudXYZI);
                            surf_filter.filter(*cloud_surf_temp_filtered);
                            cloud_surf_curr = cloud_surf_temp_filtered;
                        }

                        // 迭代优化位姿
                        for (int opt_cnt = 0; opt_cnt < 10; opt_cnt++)
                        {
                            TicToc t_opt_once;

                            // 构建优化问题
                            ceres::ResidualBlockId block_id;
                            std::vector<ceres::ResidualBlockId> residual_block_ids;

                            ceres::LossFunction *loss_function = new ceres::HuberLoss(0.1);
                            ceres::Problem::Options problem_options;
                            ceres::Problem problem(problem_options);
                            // 解决过参数化问题
                            // 四元数实际上只有三个自由度，它还受模长为1的约束
                            ceres::LocalParameterization *q_parameterization = new ceres::EigenQuaternionParameterization();
                            problem.AddParameterBlock(q_params_increm, 4, q_parameterization);
                            problem.AddParameterBlock(t_params_increm, 3);
                            // 设置速度增量约束 2m/s
                            for (int kk = 0; kk < 3; kk++)
                            {
                                problem.SetParameterLowerBound(t_params_increm, kk, -2);
                                problem.SetParameterUpperBound(t_params_increm, kk, 2);
                            }

                            // KD-Tree最近邻搜索结果的索引和距离
                            std::vector<int> pointSearchInd;
                            std::vector<float> pointSearchSqDis;

                            nearest_queue_.clear();

                            ////////// 平面点处理
                            TicToc t_surf;
                            for (int i = 0; i < cloud_surf_curr->size(); i++)
                            {

                                (*cloud_surf_curr)[i].curvature = 1; // 设置为1，等价于关闭运动补偿

                                auto p = (*cloud_surf_curr)[i];
                                Eigen::Vector3d cp(p.x, p.y, p.z);

                                Eigen::Vector3d wp = q_w_last * (Eigen::Quaterniond{1, 0, 0, 0}.slerp((*cloud_surf_curr)[i].curvature, q_w_increm) * cp + (*cloud_surf_curr)[i].curvature * t_w_increm) + t_w_last;
                                p.x = wp.x();
                                p.y = wp.y();
                                p.z = wp.z();

                                // 近邻搜索
                                ikd_tree_->Nearest_Search(p, KNN_MATCH_NUM, nearest, pointSearchSqDis);
                                nearest_queue_.push_back(nearest);

                                if (pointSearchSqDis[KNN_MATCH_NUM - 1] < 5) // 最近邻都在指定范围内
                                {
                                    // 计算质心和协方差
                                    PointCloudXYZI nearest0;
                                    for (int k = 0; k < KNN_MATCH_NUM; k++)
                                    {
                                        nearest0.push_back(nearest[k]);
                                    }
                                    Eigen::Vector4f centroid;   // 质心
                                    Eigen::Matrix3f covariance; // 协方差
                                    pcl::compute3DCentroid(nearest0, centroid);
                                    pcl::computeCovarianceMatrix(nearest0, centroid, covariance);

                                    // 计算协方差矩阵的特征值
                                    Eigen::Vector3f eigenValues;
                                    pcl::eigen33(covariance, eigenValues);

                                    std::vector<float> eigenValues_sort{eigenValues[0], eigenValues[1], eigenValues[2]};
                                    std::sort(eigenValues_sort.begin(), eigenValues_sort.end()); // 升序排列

                                    if (eigenValues_sort[1] > 3 * eigenValues_sort[0])
                                    {
                                        // 最近邻在一个平面上
                                        Eigen::Vector3d lpj{nearest0[0].x, nearest0[0].y, nearest0[0].z};
                                        Eigen::Vector3d lpl{nearest0[4].x, nearest0[4].y, nearest0[4].z};
                                        Eigen::Vector3d lpm{nearest0[7].x, nearest0[7].y, nearest0[7].z};
                                        auto cost = LidarSurfFactor::Create(cp, lpj, lpl, lpm, p.curvature, q_w_last, t_w_last);
                                        block_id = problem.AddResidualBlock(cost, nullptr, q_params_increm, t_params_increm);
                                        residual_block_ids.push_back(block_id);
                                    }
                                }
                            }

                            // ROS_INFO("surf time %f ms",  t_surf.toc());


                            TicToc t_solver;
                            ceres::Solver::Options options;
                            options.linear_solver_type = ceres::DENSE_QR; // 增量方程的解法
                            options.num_threads = 16;
                            options.max_num_iterations = 15;
                            options.minimizer_progress_to_stdout = false;
                            options.check_gradients = false;
                            options.gradient_check_relative_precision = 1e-10;
                            ceres::Solver::Summary summary; // 求解结果摘要
                            /* 求解 */
                            ceres::Solve(options, &problem, &summary);

                            // ROS_INFO("solver time %f ms",  t_solver.toc());
                            // ROS_INFO("data opt once time %f ms", t_opt_once.toc());
                            // ROS_INFO("residual_block_ids size: %ld", residual_block_ids.size());
                            // ROS_INFO("cnt size: %ld", cnt);

                            // 评估优化结果,再进行一次优化
                            if (opt_cnt > 5 && 1)
                            {

                                ceres::Problem::EvaluateOptions eval_options;
                                std::vector<double> residuals;
                                double total_cost = 0;
                                problem.Evaluate(eval_options, &total_cost, &residuals, nullptr, nullptr);
                                int cnt = 0;
                                for (int i = 0; i < residual_block_ids.size(); i++)
                                {
                                    if (fabs(residuals[i]) > 0.02) // 移除过大的残差项
                                    {
                                        // ROS_INFO("residual %d: %f", i, residuals[i]);
                                        problem.RemoveResidualBlock(residual_block_ids[i]);
                                        cnt++;
                                    }
                                }
                                // ROS_INFO("cnt: %d", cnt);
                                options.linear_solver_type = ceres::DENSE_QR; // 增量方程的解法
                                options.num_threads = 16;
                                options.max_num_iterations = 15;
                                options.minimizer_progress_to_stdout = false;
                                options.check_gradients = false;
                                options.gradient_check_relative_precision = 1e-10;
                                /* 求解 */
                                ceres::Solve(options, &problem, &summary);
                            }
                        }
                        if (1)
                        {
                            PointCloudXYZI::Ptr cloud_surf_w = boost::make_shared<PointCloudXYZI>();
                            cloud_surf_w->resize(cloud_surf_curr->size());
                            for (int i = 0; i < cloud_surf_curr->size(); i++) // 将点云投影到世界坐标系下
                            {
                                Eigen::Vector3d cp((*cloud_surf_curr)[i].x, (*cloud_surf_curr)[i].y, (*cloud_surf_curr)[i].z);
                                Eigen::Vector3d wp = q_w_last * (Eigen::Quaterniond{1, 0, 0, 0}.slerp((*cloud_surf_curr)[i].curvature, q_w_increm)
                                                                 * cp + (*cloud_surf_curr)[i].curvature * t_w_increm) + t_w_last;
                                (*cloud_surf_w)[i].x = wp.x();
                                (*cloud_surf_w)[i].y = wp.y();
                                (*cloud_surf_w)[i].z = wp.z();
                            }

                            sensor_msgs::PointCloud2 laserSurfMapMsgOut;
                            pcl::toROSMsg(*cloud_surf_w, laserSurfMapMsgOut);
                            laserSurfMapMsgOut.header.stamp = ros::Time(time_surf_cloud);
                            laserSurfMapMsgOut.header.frame_id = "map";
                            pub_surf_map.publish(laserSurfMapMsgOut);
                            
                        }
                        t_w_curr = q_w_last * t_w_increm + t_w_last;
                        q_w_curr = q_w_last * q_w_increm;

                        q_w_last = q_w_curr;
                        t_w_last = t_w_curr;

                        map_update(cloud_surf_curr);

                    }


                    // 保存轨迹到文件
                    result_fs << t_w_curr.x() << " " << t_w_curr.y() << " " << t_w_curr.z() << std::endl;

                    ROS_INFO("total time %f ms", t_total.toc());
                    ROS_INFO("-------------------------");

                    // 发布坐标转换
                    odom_trans.header.stamp = ros::Time(time_surf_cloud);
                    odom_trans.header.frame_id = "map";
                    odom_trans.child_frame_id = frame_id;
                    odom_trans.transform.translation.x = t_w_curr.x();
                    odom_trans.transform.translation.y = t_w_curr.y();
                    odom_trans.transform.translation.z = t_w_curr.z();
                    odom_trans.transform.rotation.x = q_w_curr.x();
                    odom_trans.transform.rotation.y = q_w_curr.y();
                    odom_trans.transform.rotation.z = q_w_curr.z();
                    odom_trans.transform.rotation.w = q_w_curr.w();
                    odom_broadcaster.sendTransform(odom_trans);

                    // 发布路径
                    geometry_msgs::PoseStamped pose;
                    pose.header.stamp = ros::Time(time_surf_cloud);
                    pose.header.frame_id = "map";
                    pose.pose.position.x = t_w_curr.x();
                    pose.pose.position.y = t_w_curr.y();
                    pose.pose.position.z = t_w_curr.z();
                    pose.pose.orientation.x = q_w_curr.x();
                    pose.pose.orientation.y = q_w_curr.y();
                    pose.pose.orientation.z = q_w_curr.z();
                    pose.pose.orientation.w = q_w_curr.w();
                    traj_local_msgs.poses.push_back(pose);
                    traj_local_msgs.header = pose.header;
                    pub_traj.publish(traj_local_msgs);
                }
            }
            result_fs.close();
            // pcl::io::savePCDFileASCII("/home/ghowoght/workspace/lidar_ws/map_edge.pcd", *map_edge);
            // pcl::io::savePCDFileASCII("/home/ghowoght/workspace/lidar_ws/map_surf.pcd", *map_surf);
        }
    };

};